#!/usr/bin/env python

###############################################################################
#
# MODULE:	m.fuzzy.upscale
# AUTHOR:	Thomas Leppelt
# PURPOSE:      Application of calibrated fuzzy models on raster maps with option 
#               to use empirical distribution function as input maps to generate 
#               models from.
#
###############################################################################

#%module
#% description: Application of calibrated fuzzy models on raster maps with optional use of empirical distribution function.
#% keywords: Miscellaneous
#% keywords: Fuzzy model,TAG2E, ecdf
#%end
#%option G_OPT_F_INPUT
#% description: List of input raster maps or function names for fuzzy model input. Function options: ecdf()
#% required : yes
#% multiple: yes
#% guisection: Files & format
#%end
#%option G_OPT_F_OUTPUT
#% description: Basename of output raster files.
#% required : no
#% guisection: Files & format
#%end
#%option
#% key: parameter
#% type : string
#% description: List of XML fuzzy inference parameter file names.
#% required : no
#% multiple: no
#% guisection: Files & format
#%end
#%option
#% key: file
#% type : string
#% description: File name containing XML fuzzy inference parameter file names as list.
#% required : no
#% multiple: no
#% guisection: Files & format
#%end
#%option
#% key: alias
#% type : string
#% description: Alias name of the input raster maps in the XML fuzzy inference scheme file.
#% required : yes
#% multiple: yes
#% guisection: Files & format
#%end
#%option
#% key: sample
#% type : double
#% description: Number of random samples taken from distributions.
#% required : no
#% multiple: no
#% answer: 100
#% guisection: Files & format
#%end
#%option
#% key: population
#% type : string
#% description: List of vector maps containing data used to generate specified functions. Must be provided in same order as the functions occure in input.
#% required : no
#% multiple: yes
#% guisection: Files & format
#%end
#%option
#% key: column
#% type : string
#% description: List of column names containing data used to generate specified functions. Must be provided in same order as the functions occure in input.
#% required : no
#% multiple: yes
#% guisection: Files & format
#%end
#%option
#% key: null
#% type : double
#% description: The value used for no data. Default: 9999
#% required : no
#% multiple: no
#% answer: 9999
#% guisection: Files & format
#%end
#%flag
#% key: a
#% description: Calculate additional statistics (Stdev, Var, Quantiles). Attention:  Very slow computation for big sample numbers. 
#%end
#%flag
#% key: s
#% description: Compute the cell specific standard deviation based on the rule standard deviation.
#%end
#%flag
#% key: p
#% description: Activate parallel model computing based on number of CPU cores.
#%end
"""
# Test 1: ParallelSeries.
@code

>>> test_list = []
>>> grass.run_command('g.region', n = 1000000, s = 0, e = 1000000, w = 0, overwrite = True)
0
>>> for i in range(1,1001):
...     stat = grass.run_command('r.mapcalc', expression = "test_%s = %s" %(str(i), str(i)),
...     overwrite = True, quiet = True)
...     test_list.append(str('test_' + str(i)))
...
>>> ParallelSeries(map_list = test_list, output = "test_out")
Created result map: test_out_average
>>> result = grass.raster.raster_info('test_out_average')
>>> print(result['max'])
500.5
>>> grass.run_command('g.mremove', rast="test_[0-9]*", flags = 'f', quiet = True)
0
>>> grass.run_command('g.mremove', rast="test_out_average", flags = 'f', quiet = True)
0

@endcode
"""
import random
import grass.script as grass
import multiprocessing as multi
import re
import itertools
import sys

class EmpiricalCDF:
    # Object for empirical cumulative distribution functions.
    # The expected value, or mean, of the empirical CDF is the true underlying CDF. 
    # That is, the empirical CDF is an unbiased estimate of the population CDF.
    # The empirical CDF is a consistent estimator of the true CDF at any value of x.
    # https://onlinecourses.science.psu.edu/stat464/node/84
    def __init__(self, data): # Constructor for basic ECDF properties.
        self.input_dataset = data
        self.input_sorted_set = set(data)
        self.input_count_set = [data.count(i) for i in set(data)]
        self.input_step_set = [ float(i) / len(data) for i in self.input_count_set]
        self.random_var = list()
    def random_sample(self, number): # Function for creating random samples.
        self.random_var = list()
        self.input_len = len(self.input_dataset)
        for i in range(0,number):
            self.random_var.append(self.input_dataset[random.randrange(self.input_len)])
    def __str__(self): # Print method gives list of random sample.
        return str(self.random_var)

def ParallelRun(map_list, output, function, parallel = False):
    if parallel:    
        # Get number of CPU cores.
        workers = multi.cpu_count()
    else:
        workers = 1
    # Set iteration variable for parallel processes.
    proc_count = 0
    # Set iteration variable for temporary raster maps.
    out_count = 0
    # Create empty list for model evaluations.    
    model_eval_list = []
    # Create empty model output list.
    result_list = []    
    for maps in map_list:
        # Convert input sub lists iteratively into strings.
        inmaps = ','.join(maps)
        # Create temporary output files.
        outmaps = str(output + '_' + str(out_count))
        # Evaluate list by r.series in parallel processing.
        model_eval_list.append(grass.start_command('r.series', input = inmaps, \
	                       output = outmaps, method = function, \
	                       overwrite = grass.overwrite, flags = "z"))
        # Save created raster map name to result list and print status.
        result_list.append(outmaps)
        # Increase iteration variables.
        out_count += 1
        proc_count += 1
        # Check for maximum number of parallel processes defined by cpu count.
        if (proc_count == workers) or (proc_count == int(len(map_list))):
            for process in model_eval_list:
                process.wait()
            # Empty process list and reset counter.
            model_eval_list = []
            proc_count = 0
    # Wait for remaining processes.
    for process in model_eval_list:
        process.wait()
    # Return result map list.
    return(result_list)
        
def ParallelSeries(map_list, output, function = "average", sublen = 50, parallel = False):
    # Evaluate raster maps by division into subsets with parallel processing.
    # Create empty remove list for computed maps.
    remove_temp = []
    # Create sub lists of given length. 
    chunks = [map_list[i : i + sublen] for i in range(0, len(map_list), sublen)]
    # Compute statistics for result map list with r.series.      
    result_list = ParallelRun(chunks, str(output + '_tempI'), function, parallel = parallel)
    # Add result to remove list. 
    remove_temp.append(result_list)
    # Conditional repeated parallel processing.
    if len(result_list) > int(sublen):
        chunks = [result_list[i : i + sublen] for i in range(0, len(result_list), sublen)]
        result_list = ParallelRun(chunks, str(output + '_tempII'), function)
        remove_temp.append(result_list)
    # Compute final output map.
    outmap = str(output + '_' + function) 
    inmaps = ','.join(result_list)
    grass.run_command('r.series', input = inmaps, output = outmap, method = function, \
	              overwrite = grass.overwrite, flags = "z")
    print(str("Created result map: " + output + "_" + function))
    # Format list with raster maps for removal.
    remove_temp = [",".join(l) for l in remove_temp]
    # Remove created random variable raster maps.
    for rast_list in remove_temp:
        grass.read_command('g.mremove',  rast = rast_list, flags = 'f')
        
def EvalFuzzyModelSet(map_list, parameter, alias, output_base, sample_num, sd = False, \
                        parallel = False):
    # Evaluate Fuzzy models by given raster maps or distribution functions. 
    if parallel:    
        # Get number of CPU cores.
        workers = multi.cpu_count()
    else:
        workers = 1
    # Set iteration variable for parellel processes.
    proc_count = 0
    # Set iteration variable for output raster maps.
    out_count = 0    
    # Create empty list for model evaluations.
    model_eval_list = []
    # Create empty model output list and sd output list.
    result_list = []
    sd_list = []
    # Convert map list elemets to sub lists. 
    map_list = [[i] if not isinstance(i, list) else i for i in map_list]
    # Create map combinations for fuzzy model input.
    iter_list = list(itertools.product(*map_list))
    # Loop over parameter files in parameter list. 
    for p in parameter:
        # Loop over maps in given map_list.
        for i in iter_list:
            input_maps = ','.join(list(i))
            output_name = str(output_base + "_" + str(out_count))
            
            if sd:
                sd_name = str(output_base + "_sd_" + str(out_count))
                model_eval_list.append(grass.start_command('r.fuzzy.model', \
                                    input = input_maps, \
                                    output = output_name, \
                                    parameter = p, \
                                    alias = alias, \
                                    overwrite = grass.overwrite, \
                                    sd = sd_name))
                sd_list.append(sd_name)
            else:
                model_eval_list.append(grass.start_command('r.fuzzy.model', \
                                    input = input_maps, \
                                    output = output_name, \
                                    parameter = p, \
                                    alias = alias, \
                                    overwrite = grass.overwrite))
                                    
            # Save created raster map name to result list and print status.
            result_list.append(output_name)

            print(str("Created model output: " + output_name))
            # Increase iteration variables.
            out_count += 1
            proc_count += 1
            # Check for maximum number of parallel processes defined by cpu count.
            if (proc_count == workers) or (proc_count == int(sample_num)):
                for process in model_eval_list:
                    process.wait()
                # Empty process list and reset counter.
                model_eval_list = []
                proc_count = 0
    # Wait for remaining processes.
    for process in model_eval_list:
        process.wait()
    if sd:
        return(result_list, sd_list)
    else:
        return(result_list)
        
def main():
  
    input = options['input']
    output = options['output']
    parameter = options['parameter']
    alias = options['alias']
    sample = options['sample']
    population = options['population']
    column = options['column']
    filename = options['file']
    navalue = options['null']
    statistics = flags['a']  
    sd = flags['s']
    parallel = flags['p']
    # List of optional input functions to calculate random values. 
    function_list = ["ecdf()", "cdf()", "norm()"]
    
    # Check if parameters or file list is given.
    try:
        file_list = open(str(filename))
        d = file_list.read()
        parameter_list = d.split()
        file_list.close()
    except:
        parameter_list = parameter.split(',')
        if parameter_list == [""]:
            grass.fatal(("Parameter files not given."))

    # Format input parameter string.
    input_parameter = input.split(',')
    input_parameter = [i.strip() for i in input_parameter]

    # Check if input raster maps exists in actual mapset.
    input_list = [i for i in input_parameter if i not in function_list]
    for map in input_list:
        if grass.find_file(map)['file'] == "":
            grass.fatal(("Raster map <%s> not exists.") % map)
            
    # Format input population string.
    popu_input = population.split(',')
    popu_input = [i.strip() for i in popu_input]

    # Format input population string.
    col_input = column.split(',')
    col_input = [i.strip() for i in col_input]
    
    # Create list of boolean expressions for functions input list.
    function_bool = [[i == j for j in function_list] for i in input_parameter]

    # Loop through input parameters and evaluate given function statements and 
    # return raster map list.
    
    # Create empty list for input maps.
    map_list = list()
    # Create empty list for maps that will be removed after computations.
    remove_list = list()
    # Set iteration variable for population number. 
    popu_index = 0
    for i in range(0, len(input_parameter)):
        if True in function_bool[i]:
	    # Get function index and type.
	    function_index = function_bool[i].index(True)
	    function_type = function_list[function_index]	    
	    # Create list from vector map column with population data.
	    popu_table = grass.read_command('v.db.select', \
	                 map = popu_input[popu_index], \
	                 columns = col_input[popu_index], \
	                 flags = 'c', sep= ';')    
	    # Format population data.
	    popu_table = re.split('\n', popu_table.strip('\n'))	    
	    # Remove NA values from population data.
	    popu_table = [item for item in popu_table if item != str(navalue)]	    
	    # Create random values for selected function of given population.
	    if function_type == "ecdf()":
	        popu_function = EmpiricalCDF(popu_table)
                popu_function.random_sample(int(sample))
            	popu_sample = popu_function.random_var         
            # Create raster maps based on sampled random variable.
            raster_list = []
            # Loop through number of sample random variables. 
            for i in range(0, int(sample)):
	        raster_name = str(popu_input[popu_index] + '_' + \
	                      col_input[popu_index] + '_' + str(i))
                # Create raster maps with random variables sampled from given
                # distribution function. 
                grass.read_command('r.mapcalc', \
	                           expression = str(raster_name + '=' + \
	                           popu_sample[i]), overwrite = grass.overwrite)
	        raster_list.append(raster_name)
	        print(str("Created sample raster map: " + raster_name))
	    # Add raster map names to result lists. 
	    map_list.append(raster_list)
	    remove_list.append(raster_list)
	    popu_index += 1
        else:
	    map_list.append(input_parameter[i])

    # Use raster map list to compute fuzzy models.
    if sd:
        result_list, sd_list = EvalFuzzyModelSet(map_list, parameter_list, \
                                                    alias, output, sample, sd, \
                                                    parallel)
    else:
        result_list = EvalFuzzyModelSet(map_list, parameter_list, alias, output, \
                                        sample, sd, parallel)

    # Compute statistics for result map list with r.series.      
    # Average output raster map.
    ParallelSeries(result_list, output, "average", parallel = parallel)
    ParallelSeries(result_list, output, "maximum", parallel = parallel)
    ParallelSeries(result_list, output, "minimum", parallel = parallel)
    
    if sd:
        ParallelSeries(sd_list, output + "_sd", "average", parallel = parallel)
        ParallelSeries(sd_list, output + "_sd", "maximum", parallel = parallel)
        ParallelSeries(sd_list, output + "_sd", "minimum", parallel = parallel)
        
    # Additional statistics.	       
    if statistics:
        # Format result raster map list.
        result_str = '\n'.join(result_list)
        # Create new file in temporary directory with result list as input for r.series.
        result_path = "/tmp/m_fuzzy_upscale_result_list.txt"
        result_file = open(result_path, "w")
        result_file.write(result_str)
        result_file.close()
        # Standard deviation output raster map.
        grass.read_command('r.series', file = result_path, \
	                   output = str(output + "_stdev"), method = "stddev", \
	                   overwrite = grass.overwrite, flags = "z")
        print(str("Created result map: " + output + "_stdev"))
        # Variance output raster map.
        grass.read_command('r.series', file = result_path, \
	                   output = str(output + "_var"), method = "variance", \
	                   overwrite = grass.overwrite, flags = "z")	                 
        print(str("Created result map: " + output + "_var"))
        # Q05 output raster map.
        grass.read_command('r.series', file = result_path, quantile = "0.05", \
	                   output = str(output + "_q05"), method = "quantile", \
	                   overwrite = grass.overwrite, flags = "z")
        print(str("Created result map: " + output + "_q05"))
        # Q95 output raster map.
        grass.read_command('r.series', file = result_path, quantile = "0.95", \
	                   output = str(output + "_q95"), method = "quantile", \
	                   overwrite = grass.overwrite, flags = "z")	                 
        print(str("Created result map: " + output + "_q95"))
        # Q50 output raster map.
        grass.read_command('r.series', file = result_path, quantile = "0.5", \
	                   output = str(output + "_q50"), method = "quantile", \
	                   overwrite = grass.overwrite, flags = "z")
	print(str("Created result map: " + output + "_q50"))

    # Format list with raster maps for removal.
    remove_list = [",".join(l) for l in remove_list]
    # Remove fuzzy model raster maps.
    grass.read_command('g.mremove',  rast = str(output + '_' + '[0-9]'), flags = 'rf')
    # Remove created random variable raster maps.
    for rast_list in remove_list:
        grass.read_command('g.mremove',  rast = rast_list, flags = 'f')
     
if __name__ == "__main__":
    if len(sys.argv) == 2 and sys.argv[1] == 'doctest':
        import doctest
        doctest.testmod()
        exit()
    options, flags = grass.parser()
    main()
